#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import asyncio
import fractions
import io
import logging
import socket
import threading
import time
from typing import Optional

from common.config import ServiceConfig
import numpy as np
from PIL import Image, ImageDraw, ImageFont
from aiortc import VideoStreamTrack
from av import VideoFrame

logger = logging.getLogger(__name__)

class WebRTCVideoTrack(VideoStreamTrack):
    """WebRTC视频轨道"""
    def __init__(self):
        super().__init__()
        self.__kind = "video"
        self.__connected = False
        self.__frame_socket = None
        self.__buffer = bytearray()
        self.__last_jpeg_bytes = None
        self.__blank_frame = None  # 缓存空白帧
        self.__connect_task = None

        # 时间戳基准（90kHz时钟）
        self._ts_base = time.perf_counter()

        # 启动异步连接任务
        # self.__connect_task = asyncio.create_task(self.__connect_device())
        self.__connect_thread = self.__connect_remote_server()
    
    def __del__(self):
        self.destroy()
    
    async def __connect_device(self):
        """异步连接设备"""
        self.__frame_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        # 设置非阻塞模式
        self.__frame_socket.setblocking(False)
        try:
            # 使用异步连接
            loop = asyncio.get_event_loop()
            await loop.sock_connect(self.__frame_socket, ('127.0.0.1', ServiceConfig.DEVICE_FRAME_PORT))
            self.__connected = True
            logger.info("视频轨道连接设备成功")
        except Exception as e:
            logger.error(f"视频轨道连接设备失败: {e}")
            if self.__frame_socket:
                self.__frame_socket.close()
                self.__frame_socket = None

    def __connect_remote_server(self):
        """创建线程同步连接远程服务器"""
        def connect_thread():
            """连接线程函数"""
            try:
                # 创建新的socket
                frame_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
                # 同步连接
                frame_socket.connect(('0.0.0.0', ServiceConfig.DEVICE_FRAME_PORT))
                
                # 连接成功后更新实例变量
                if self.__frame_socket:
                    self.__frame_socket.close()
                self.__frame_socket = frame_socket
                self.__connected = True
                logger.info("视频轨道线程连接设备成功")
            except Exception as e:
                logger.error(f"视频轨道线程连接设备失败: {e}")
                self.__connected = False
        
        # 创建并启动连接线程
        connect_thread_obj = threading.Thread(target=connect_thread, daemon=True)
        connect_thread_obj.start()
        return connect_thread_obj

    def destroy(self):
        if self.__connect_task:
            self.__connect_task.cancel()
        if self.__frame_socket:
            self.__frame_socket.close()
            self.__frame_socket = None

    async def recv(self) -> VideoFrame:
        """接收视频帧"""
        if self.__connected is False:
            logger.warning("视频轨道未连接设备，无法接收视频帧")
            return self._create_blank_frame()
        try:
            # 接收JPEG帧数据
            frame_data = await self._recv_jpeg_frame()
            if frame_data is None:
                # 如果没有收到完整帧，返回上一帧（若存在），否则返回空白帧
                if self.__last_jpeg_bytes is not None:
                    return self._jpeg_to_video_frame(self.__last_jpeg_bytes)
                return self._create_blank_frame()
            
            # 缓存上一帧的JPEG数据
            self.__last_jpeg_bytes = frame_data
            
            # 将JPEG数据转换为VideoFrame
            return self._jpeg_to_video_frame(frame_data)
            
        except Exception as e:
            logger.error(f"接收视频帧失败: {e}")
            # 异常时尽量返回上一帧（若存在），否则返回空白帧
            try:
                if self.__last_jpeg_bytes:
                    return self._jpeg_to_video_frame(self.__last_jpeg_bytes)
            except Exception:
                pass
            return self._create_blank_frame()
    
    async def _recv_jpeg_frame(self) -> Optional[bytes]:
        """从socket接收完整的JPEG帧数据"""
        if not self.__frame_socket:
            return None
        try:
            # 使用异步方式从非阻塞socket读取数据
            loop = asyncio.get_event_loop()
            chunk = await loop.sock_recv(self.__frame_socket, 131072)
            if not chunk:
                # 连接已关闭
                self.__connected = False
                try:
                    self.__frame_socket.close()
                finally:
                    self.__frame_socket = None
                return None

            self.__buffer.extend(chunk)
            # 检查是否有完整的JPEG帧
            if len(self.__buffer) >= 4:
                # 查找JPEG开始标记 (0xFFD8)
                start_idx = self.__buffer.find(b'\xff\xd8')
                if start_idx != -1:
                    # 查找JPEG结束标记 (0xFFD9)
                    end_idx = self.__buffer.find(b'\xff\xd9', start_idx + 2)
                    if end_idx != -1:
                        # 提取完整的JPEG帧
                        frame_data = bytes(self.__buffer[start_idx:end_idx + 2])
                        # 移除已处理的字节
                        del self.__buffer[:end_idx + 2]
                        return frame_data
        except (BlockingIOError, InterruptedError):
            # 暂无可读数据
            return None
        except Exception as e:
            logger.error(f"读取JPEG帧数据失败: {e}")
            return None

        return None

    
    def _jpeg_to_video_frame(self, jpeg_data: bytes) -> VideoFrame:
        """将JPEG数据转换为VideoFrame"""
        try:
            # 使用PIL解码JPEG
            pil_image = Image.open(io.BytesIO(jpeg_data))
            
            # 转换为RGB格式
            if pil_image.mode != 'RGB':
                pil_image = pil_image.convert('RGB')
            
            # 转换为numpy数组
            frame_array = np.array(pil_image)
            
            # 创建VideoFrame
            video_frame = VideoFrame.from_ndarray(frame_array, format="rgb24")
            
            # 设置时间戳
            pts, time_base = self.next_timestamp()
            video_frame.pts = pts
            video_frame.time_base = time_base
            
            return video_frame
            
        except Exception as e:
            logger.error(f"JPEG转VideoFrame失败: {e}")
            return self._create_blank_frame()
    
    def _create_blank_frame(self) -> VideoFrame:
        """创建空白帧（使用缓存优化）"""
        if self.__blank_frame is None:
            # 首次创建空白帧并缓存
            frame_array = np.zeros((480, 640, 3), dtype=np.uint8)
            
            # 添加"No Signal"文字
            try:
                pil_image = Image.fromarray(frame_array)
                draw = ImageDraw.Draw(pil_image)
                
                # 尝试使用默认字体
                try:
                    font = ImageFont.load_default()
                except:
                    font = None
                font = self._get_font(size=56)
                 
                # 在中心绘制"No Signal"文字
                text = "No Signal"
                if font:
                    bbox = draw.textbbox((0, 0), text, font=font)
                    text_width = bbox[2] - bbox[0]
                    text_height = bbox[3] - bbox[1]
                else:
                    # 无法加载大字体时的估算（尽量大一些）
                    text_width = len(text) * 28
                    text_height = 56
                
                x = (640 - text_width) // 2
                y = (480 - text_height) // 2
                
                draw.text((x, y), text, fill=(255, 255, 255), font=font)
                
                frame_array = np.array(pil_image)
                
            except Exception as e:
                logger.error(f"绘制空白帧文字失败: {e}")
            
            # 缓存基础帧数组
            self.__blank_frame = frame_array
        
        # 每次都创建新的VideoFrame（因为需要不同的时间戳）
        video_frame = VideoFrame.from_ndarray(self.__blank_frame.copy(), format="rgb24")
        
        # 设置时间戳
        pts, time_base = self.next_timestamp()
        video_frame.pts = pts
        video_frame.time_base = time_base
        
        return video_frame
    
    def next_timestamp(self):
        """返回 (pts, time_base)，基于90kHz时钟生成递增时间戳"""
        if not hasattr(self, "_ts_base") or self._ts_base is None:
            self._ts_base = time.perf_counter()
        elapsed = time.perf_counter() - self._ts_base
        return int(elapsed * 90000), fractions.Fraction(1, 90000)

    def _get_font(self, size: int = 56):
        """尝试加载常见的TrueType字体，失败则回退到默认字体"""
        font_paths = [
            "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf",
            "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf",
            "/usr/share/fonts/truetype/liberation/LiberationSans-Bold.ttf",
            "/usr/share/fonts/truetype/liberation/LiberationSans-Regular.ttf",
            "/usr/share/fonts/truetype/freefont/FreeSansBold.ttf",
            "/usr/share/fonts/truetype/freefont/FreeSans.ttf",
            "arial.ttf",
        ]
        for path in font_paths:
            try:
                return ImageFont.truetype(path, size)
            except Exception:
                continue
        try:
            return ImageFont.load_default()
        except Exception:
            return None
        
